#include "ip_packet.h"
#include <iostream>
#include <fstream>
#include <cstring>

// 封装函数，将数据封装到IP数据包并写入文件
void encapsulate(const std::string& data, unsigned int sourceIP, unsigned int destIP, const std::string& filename) {
    // 限制数据的最大长度
    if (data.size() > MAX_DATA_SIZE) {
        std::cerr << "数据超过最大长度 " << MAX_DATA_SIZE << " 字节，无法封装。" << std::endl;
        return;
    }

    IPv4Header header;
    memset(&header, 0, sizeof(IPv4Header));

    // 设置IPv4头部字段
    header.version = 4; // IPv4
    header.ihl = 5;     // 头部长度为5*4=20字节
    header.tos = 0;     // 服务类型
    header.totalLength = sizeof(IPv4Header) + data.size();
    header.identification = 0x1234; // 随机标识符
    header.flags = 2;   // 不分段标志
    header.fragmentOffset = 0;
    header.ttl = 64;    // 默认生存时间
    header.protocol = 6; // 假设为TCP协议
    header.sourceIP = sourceIP;
    header.destIP = destIP;

    // 计算校验和
    header.checksum = 0;
    header.checksum = calculateChecksum((unsigned short*)&header, sizeof(IPv4Header));

    // 创建数据包（头部 + 数据）
    unsigned char packet[sizeof(IPv4Header) + MAX_DATA_SIZE];  // 固定大小的数组
    memset(packet, 0, sizeof(packet));

    // 拷贝头部到数据包
    memcpy(packet, &header, sizeof(IPv4Header));

    // 拷贝数据到数据包（如果数据大小超过最大限制，只拷贝最大字节数）
    memcpy(packet + sizeof(IPv4Header), data.c_str(), data.size());

    // 将数据包写入文件
    std::ofstream outFile(filename, std::ios::binary);
    if (outFile.is_open()) {
        outFile.write(reinterpret_cast<const char*>(packet), sizeof(IPv4Header) + data.size());
        outFile.close();
        std::cout << "数据已成功封装并写入文件: " << filename << std::endl;
    } else {
        std::cerr << "文件打开失败！" << std::endl;
    }
}

int main() {
    std::string data = "Hello, this is a test message!";
    unsigned int sourceIP = 0xC0A80001; // 192.168.0.1
    unsigned int destIP = 0xC0A80002;   // 192.168.0.2

    // 调用封装函数，数据将写入 "ip_packet.dat" 文件
    encapsulate(data, sourceIP, destIP, "ip_packet.dat");

    return 0;
}
